import json

from openai import OpenAI

import os
import requests

from django.conf import settings

MODEL_API_HOST = "http://" + settings.MODEL_API_HOST + ":8010/v1/"


class ModelAPI:
    def __init__(self):
        """
        初始化 API 客户端。
        """
        # 清除可能存在的代理环境变量
        os.environ['http_proxy'] = ''
        os.environ['https_proxy'] = ''
        os.environ['no_proxy'] = '*'
        self.client = OpenAI(api_key="EMPTY", base_url=MODEL_API_HOST)

    def chat_stream(self, messages, model='chatglm3-6b', max_tokens=6000, temperature=0.1,
                    presence_penalty=0.1, top_p=1):
        """
        发送请求到模型 API 并处理响应。
        :param model: 使用的模型名称。
        :param messages: 消息列表，每个消息是一个包含角色和内容的字典。
        :param use_stream: 是否使用流式响应。
        :param max_tokens: 最大令牌数。
        :param temperature: 温度参数。
        :param presence_penalty: 出现罚分。
        :param top_p: 概率阈值。
        """

        real_messages = [
            # {"role": "system",
            #  "content": "You are ChatGLM3, a large language model trained by Zhipu.AI. Follow the user's instructions carefully. Respond using markdown."},
            {"role": "user", "content": messages}
        ]

        response = self.client.chat.completions.create(
            model=model,
            messages=real_messages,
            stream=True,
            max_tokens=max_tokens,
            temperature=temperature,
            # presence_penalty=presence_penalty,
            top_p=top_p
        )

        if response:
            for chunk in response:
                yield json.dumps({'text': chunk.choices[0].delta.content}) + "&&##&&"
        else:
            print("Error:", response.status_code)

    def chat(self, messages, model='chatglm3-6b', max_tokens=1024, temperature=0.1,
             presence_penalty=0.1, top_p=1):
        """
        发送请求到模型 API 并处理响应。
        :param model: 使用的模型名称。
        :param messages: 消息列表，每个消息是一个包含角色和内容的字典。
        :param use_stream: 是否使用流式响应。
        :param max_tokens: 最大令牌数。
        :param temperature: 温度参数。
        :param presence_penalty: 出现罚分。
        :param top_p: 概率阈值。
        """

        real_messages = [
            # {"role": "system",
            #  "content": "You are ChatGLM3, a large language model trained by Zhipu.AI. Follow the user's instructions carefully. Respond using markdown."},
            {"role": "user", "content": messages}
        ]

        response = self.client.chat.completions.create(
            model=model,
            messages=real_messages,
            stream=False,
            max_tokens=max_tokens,
            temperature=temperature,
            # presence_penalty=presence_penalty,
            top_p=top_p
        )

        if response:
            content = response.choices[0].message.content
            print(content)
            return content
        else:
            print("Error:", response.status_code)


# 使用示例
if __name__ == "__main__":
    # 假设 client 是正确配置的 OpenAI 客户端
    api = ModelAPI()
    messages = [
        # {"role": "system",
        #  "content": "You are ChatGLM3, a large language model trained by Zhipu.AI. Follow the user's instructions carefully. Respond using markdown."},
        {"role": "user", "content": ""}
    ]
    api.chat("你好")
